"""
python fujian1_statsArima_after_smooth.py
"""


import os
import json
import pandas as pd
from statsmodels.tsa.arima.model import ARIMA
from pmdarima import auto_arima

# 定义输入和输出路径
input_dir = "./fujian/fujian1/spline_then_smooth"
output_dir = "./fujian/fujian1/smooth_then_statsArima"

# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)

# 列出并打印 input_dir 中的文件个数和文件名
files = [f for f in os.listdir(input_dir) if f.endswith(".json")]
print(f"在目录 {input_dir} 中找到 {len(files)} 个文件：")
for f in files:
    print(f)

# 处理每一个 JSON 文件
for filename in files:
    input_path = os.path.join(input_dir, filename)
    
    # 读取 JSON 文件
    with open(input_path, 'r') as f:
        data = json.load(f)

    # 转换为 DataFrame
    df = pd.DataFrame(data)
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)

    # 提取库存数据
    inventory_data = df['inventory']

    # 使用 auto_arima 自动选择合适的 ARIMA 参数
    model = auto_arima(inventory_data, seasonal=False, stepwise=True, trace=True)

    # 预测未来 3 个月的数据
    forecast = model.predict(n_periods=3)

    # 创建包含原始数据和预测数据的新 DataFrame
    future_dates = pd.date_range(start='2023-07-01', periods=3, freq='MS')
    forecast_df = pd.DataFrame({'date': future_dates, 'inventory': forecast})

    # 合并原始数据和预测数据
    combined_data = pd.concat([df.reset_index(), forecast_df], ignore_index=True)

    # 转换为 JSON 格式
    output_json = combined_data.to_json(orient='records', date_format='iso')
    output_data = json.loads(output_json)

    # 输出到新的 JSON 文件
    category_number = filename.split('category')[1].split('.')[0]
    output_filename = f"forecastAfterSpline_category{category_number}.json"
    output_path = os.path.join(output_dir, output_filename)

    with open(output_path, 'w') as f:
        json.dump(output_data, f, indent=4)

    print(f"{filename} 的预测数据已经保存为 {output_filename}")

print("预测处理完成，文件已输出到:", output_dir)
